import evaluate
import torch
from datasets import load_dataset
from transformers import DataCollatorWithPadding, AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from transformers import pipeline
torch.cuda.empty_cache()
torch.cuda.set_device(0)

# 加载数据集
# load_dataset: 加载数据集函数
# "csv": 指定数据集格式为CSV
# data_files = "./ChnSentiCorp_htl_all.csv": 指定CSV文件路径
# split = "train": 指定加载训练集分割
# "train": 训练集，通常用于训练模型。
# "test": 测试集，用于评估模型的最终性能。
# "validation" 或 "valid": 验证集，用于调整超参数和防止过拟合。
# "all": 加载整个数据集，不考虑任何预定义的分割。
dataset = load_dataset("csv", data_files = "./ChnSentiCorp_htl_all.csv", split = "train")
# 数据清洗：去除review字段为空的样本
# filter: 过滤函数
# lambda x: x["review"] is not None: 过滤条件，保留review不为空的样本
dataset = dataset.filter(lambda x: x["review"] is not None)
print("------------------------------------------1---------------------------------------------")
# 划分数据集
# train_test_split: 将数据集划分为训练集和测试集
# test_size = 0.1: 测试集占总数据集的10%
datasets = dataset.train_test_split(test_size = 0.1)
print("------------------------------------------2---------------------------------------------")
# 数据集预处理
# AutoTokenizer.from_pretrained("hfl/rbt3"): 加载预训练的中文分词器
# "hfl/rbt3": 预训练模型名称，是哈工大讯飞实验室的中文BERT模型
tokenizer = AutoTokenizer.from_pretrained("hfl/rbt3")

# 定义数据预处理函数
def process_function(examples):
	# tokenizer(examples["review"], max_length = 128, truncation = True): 对文本进行分词
	# examples["review"]: 输入的评论文本
	# max_length = 128: 设置最大序列长度为128个token
	# truncation = True: 当文本超过最大长度时进行截断
	tokenized_examples = tokenizer(examples["review"], max_length = 128, truncation = True)
	# 将原始标签赋值给模型需要的labels字段
	# examples["label"]: 原始标签数据
	tokenized_examples["labels"] = examples["label"]
	return tokenized_examples

# 应用数据预处理Transformers 库期望的模型输入格式
# datasets.map(): 对数据集应用处理函数
# process_function: 处理函数
# batched = True: 批量处理数据以提高效率
# remove_columns = datasets["train"].column_names: 移除原始列，只保留处理后的数据
tokenized_datasets = datasets.map(process_function, batched = True, remove_columns = datasets["train"].column_names)
print("------------------------------------------3---------------------------------------------")
# 创建模型
# AutoModelForSequenceClassification.from_pretrained("hfl/rbt3"): 加载预训练的序列分类模型
# "hfl/rbt3": 预训练模型名称
model = AutoModelForSequenceClassification.from_pretrained("hfl/rbt3")


# 定义确保权重连续性的函数
def ensure_weights_contiguous(model):
	# model.named_parameters(): 获取模型中所有命名参数
	for name, param in model.named_parameters():
		# param.is_contiguous(): 检查参数在内存中是否连续存储
		if not param.is_contiguous():
			print(f"Making {name} contiguous.")
			# param.data.contiguous(): 创建连续存储版本并替换原数据
			param.data = param.data.contiguous()

# 调用函数确保模型权重连续性
ensure_weights_contiguous(model)
print("------------------------------------------4----------------------------------------------")
# 创建评估函数
# evaluate.load("accuracy"): 加载准确率评估指标
acc_metric = evaluate.load("accuracy")
# evaluate.load("f1"): 加载F1分数评估指标
f1_metric = evaluate.load("f1")


def eval_metric(eval_predict):
	# eval_predict: 包含预测结果和真实标签的元组
	predictions, labels = eval_predict
	# predictions.argmax(axis = -1): 获取预测概率最高的类别
	# axis = -1: 沿着最后一个维度进行argmax操作
	predictions = predictions.argmax(axis = -1)
	# 计算准确率
	# predictions: 预测结果
	# references = labels: 真实标签
	acc = acc_metric.compute(predictions = predictions, references = labels)
	# 计算F1分数
	f1 = f1_metric.compute(predictions = predictions, references = labels)
	# acc.update(f1): 将F1分数合并到准确率结果中
	acc.update(f1)
	return acc
print("------------------------------------------5----------------------------------------------")
# 创建TrainingArguments
# 创建训练参数
# TrainingArguments: 训练参数配置类
train_args = TrainingArguments(
	output_dir = "./checkpoints",              # 输出文件夹，保存模型和日志
	per_device_train_batch_size = 128,         # 训练时每个设备的batch_size
	per_device_eval_batch_size = 256,          # 验证时每个设备的batch_size
	fp16 = True,                               # 启用混合精度训练，节省内存并加速训练
	logging_steps = 10,                        # 每10步打印一次训练日志
	eval_strategy="epoch",                     # 每个epoch进行一次评估
	save_strategy = "epoch",                   # 每个epoch保存一次模型
	save_total_limit = 3,                      # 最多保存3个模型检查点
	learning_rate = 2e-5,                      # 学习率设置为2e-5
	weight_decay = 0.01,                       # 权重衰减系数，用于正则化防止过拟合
	metric_for_best_model = "f1",              # 以F1分数作为选择最佳模型的指标
	load_best_model_at_end = True              # 训练完成后加载最佳模型
)
print("------------------------------------------6----------------------------------------------")
# 创建Trainer
# Trainer: 训练器类，用于模型训练和评估
trainer = Trainer(
	model = model,                                    # 要训练的模型
	args = train_args,                                # 训练参数
	train_dataset = tokenized_datasets["train"],      # 训练数据集
	eval_dataset = tokenized_datasets["test"],        # 验证数据集
	data_collator = DataCollatorWithPadding(tokenizer = tokenizer),  # 数据整理器，用于动态填充批次
	compute_metrics = eval_metric                     # 评估指标计算函数
)
print("------------------------------------------7----------------------------------------------")
# 模型训练
trainer.train()
print("------------------------------------------8----------------------------------------------")
# 模型评估
# trainer.evaluate(tokenized_datasets["test"]): 在测试集上评估模型
trainer.evaluate(tokenized_datasets["test"])
print("------------------------------------------9----------------------------------------------")
# 模型预测
# trainer.predict(tokenized_datasets["test"]): 在测试集上进行预测
trainer.predict(tokenized_datasets["test"])
# 设置标签映射
# {0: "差评！", 1: "好评！"}: 将数字标签映射为中文标签
id2_label = id2_label = {0: "差评！", 1: "好评！"}
# 将标签映射配置到模型中
model.config.id2label = id2_label
pipe = pipeline("text-classification", model = model, tokenizer = tokenizer, device = 0)
sen = "我觉得不错！"
print(pipe(sen))
print("------------------------------------------10----------------------------------------------")